"""Run a single task from the tasks pickle with a reduced number of rollouts.

Usage:
  python tools/run_one_task.py <config_path> [--rollouts N] [--task-id ID]

This script loads the same config objects used by `mcts_runner`, overrides
`max_rollout_steps` in-memory, picks either the provided task id or the first
task in the pickle, and calls `MCTSRunner.run_one_task` (single-process) so you
can debug/log in the foreground.
"""
import sys
import argparse
import pickle
from pathlib import Path
from alphasql.runner.mcts_runner import MCTSRunner


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("config", help="Path to MCTS config (yaml/json)")
    parser.add_argument("--rollouts", type=int, default=10, help="Override max_rollout_steps to this small value")
    parser.add_argument("--task-id", type=int, default=None, help="Question id to run (if not provided, uses first task)")
    args = parser.parse_args()

    runner = MCTSRunner(args.config)

    # Override rollout steps
    print(f"Original max_rollout_steps: {runner.config.max_rollout_steps}")
    runner.config.max_rollout_steps = args.rollouts
    print(f"Overriding max_rollout_steps -> {runner.config.max_rollout_steps}")

    tasks_path = Path(runner.config.tasks_file_path)
    assert tasks_path.exists(), f"Tasks pickle not found: {tasks_path}"

    with open(tasks_path, "rb") as f:
        tasks = pickle.load(f)

    if args.task_id is None:
        task = tasks[0]
        print(f"No task id provided, running first task with id: {task.question_id}")
    else:
        task = next((t for t in tasks if t.question_id == args.task_id), None)
        if task is None:
            raise ValueError(f"Task id {args.task_id} not found in tasks pickle")
        print(f"Running task id: {task.question_id}")

    # Run in-process to make debugging easier
    runner.run_one_task(task)


if __name__ == "__main__":
    main()
